package com.leozhang.ElasticsearchPagination;

import org.apache.http.HttpHost;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class ElasticsearchPagination {
    private final RestHighLevelClient client;
    private int currentPage = 1; // 当前页码
    private final int pageSize = 10; // 每页大小
    private final List<Object[]> sortValues; // 缓存每页起点的 sort 值

    public ElasticsearchPagination() {
        // 初始化 Elasticsearch 客户端
        client = new RestHighLevelClient(
                RestClient.builder(new HttpHost("10.40.0.9", 9200, "http")));
        sortValues = new ArrayList<>();
    }

    // 关闭客户端
    public void close() throws IOException {
        client.close();
    }

    // 构建查询
    private SearchRequest buildSearchRequest(Object[] searchAfter) {
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
                .query(QueryBuilders.matchAllQuery()) // 示例查询，可替换为你的查询条件
                .from(0)
                .size(pageSize)
                .sort("suricata.eve.timestamp", SortOrder.DESC) // 按 timestamp 降序
                .sort("_id", SortOrder.ASC);

        if (searchAfter != null) {
            sourceBuilder.searchAfter(searchAfter);
        }

        return new SearchRequest("filebeat-6.8.0-2025.04.03").source(sourceBuilder); // 替换为你的索引名
    }

    // 获取第一页
    public List<SearchHit> getFirstPage() throws IOException {
        SearchRequest request = buildSearchRequest(null);
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);
        SearchHit[] hits = response.getHits().getHits();

        if (hits.length > 0) {
            sortValues.clear(); // 重置缓存
            sortValues.add(hits[0].getSortValues()); // 缓存第 1 页起点
            currentPage = 1;
        }

        return Arrays.asList(hits);
    }

    // 下一页
    public List<SearchHit> nextPage() throws IOException {
        if (currentPage <= 0) {
            return getFirstPage(); // 如果未初始化，从第一页开始
        }

        // 使用当前页最后一条的 sort 值
        SearchRequest request = buildSearchRequest(getLastSortValue());
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);
        SearchHit[] hits = response.getHits().getHits();

        if (hits.length > 0) {
            currentPage++;
            sortValues.add(hits[0].getSortValues()); // 缓存新页起点
            return Arrays.asList(hits);
        } else {
            System.out.println("已经是最后一页");
            return new ArrayList<>();
        }
    }

    // 上一页
    public List<SearchHit> prevPage() throws IOException {
        if (currentPage <= 1) {
            System.out.println("已经是第一页");
            return new ArrayList<>();
        }

        // 使用前一页的 sort 值
        Object[] prevSort = sortValues.get(currentPage - 2); // 前一页起点
        SearchRequest request = buildSearchRequest(prevSort);
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);
        SearchHit[] hits = response.getHits().getHits();

        if (hits.length > 0) {
            currentPage--;
            return Arrays.asList(hits);
        }

        return new ArrayList<>();
    }

    // 获取当前页最后一条的 sort 值（用于下一页）
    private Object[] getLastSortValue() throws IOException {
        SearchRequest request = buildSearchRequest(currentPage == 1 ? null : sortValues.get(currentPage - 1));
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);
        SearchHit[] hits = response.getHits().getHits();
        return hits.length > 0 ? hits[hits.length - 1].getSortValues() : null;
    }

    // 获取当前页码
    public int getCurrentPage() {
        return currentPage;
    }

    // 测试代码
    public static void main(String[] args) throws IOException {
        ElasticsearchPagination pagination = new ElasticsearchPagination();

        // 获取第 1 页
        System.out.println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> Page 1:");
        List<SearchHit> hits = pagination.getFirstPage();
        printHits(hits);

        // 下一页
        for (int i = 0; i < 10; i++) {
            System.out.println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> Page 2:");
            hits = pagination.nextPage();
            printHits(hits);

            // 再下一页
            // System.out.println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> Page 3:");
            // hits = pagination.nextPage();
            // printHits(hits);
        }

        // 上一页
        // for (int i = 0; i < 10; i++) {
        //     System.out.println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> Back to Page 2:");
        //     hits = pagination.prevPage();
        //     printHits(hits);

        //     // 再上一页
        //     System.out.println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> Back to Page 1:");
        //     hits = pagination.prevPage();
        //     printHits(hits);
        // }

        pagination.close();

        pagination.printSortValues();
    }

    private void printSortValues() {
        System.out.println("Sort value length: " + sortValues.size());
        sortValues.stream().map(Arrays::toString).forEach(System.out::println);
    }

    private static void printHits(List<SearchHit> hits) {
        for (SearchHit hit : hits) {
            System.out.println(hit.getSourceAsString());
        }
        System.out.println("Total hits: " + hits.size());
    }
}
